Flash Attention

Attention

• The softmax function is a key component of the attention mechanism.

• It is used to compute the attention weights, which scale the values based on the similarity of the queries and keys. \[ \text{Attention}(Q, K, V) = \text{Softmax}(\frac{QK^\mathsf{T}}{\sqrt{d_k}}) V, \text{where} \\ Q, K, V \in \mathbb{R}^{N \times d_k}, \\ N \text{ is the sequence length,}\\ d_k \text{ is the dimension of the heads.} \]

Safe softmax

• In order to mitigate this the followint trick is used: \[ \text{Softmax}(x)_i = \frac{e^{x_i}}{\sum_{j=1}^{N} e^{x_j}} = \\ \frac{C}{C}\frac{e^{x_i}}{\sum_{j=1}^{N} e^{x_j}} = \\ \frac{C e^{x_i}}{\sum_{j=1}^{N} C e^{x_j}} = \frac{e^{x_i + log(C)}}{\sum_{j=1}^{N} e^{x_j + log(C)}},\\ \text{where typically } log(C) = -\max_{i} x_i \]

Safe softmax

3 pass implementation

• The naive implementation requires three passes over the data.

• Lets introduce the following notation:

  • \(\left\{m_{i}\right\}: \max _{j=1}^{i}\left\{x_{j}\right\}\), with initial value \(m_{0}=-\infty\).
  • \(\left\{d_{i}\right\}: \sum_{j=1}^{i} e^{x_{j}-m_{N}}\), with initial value \(d_{0}=0\).
  • \(\left\{a_{i}\right\}\) : the final softmax value.

Safe softmax

3 pass implementation

1. For \(i \leftarrow 1, N\) do \[ m_{i} \leftarrow \max \left(m_{i-1}, x_{i}\right) \qquad(1)\]

2. For \(i \leftarrow 1, N\) do \[ d_{i} \leftarrow d_{i-1}+e^{x_{i}-m_{N}} \qquad(2)\]

3. For \(i \leftarrow 1, N\) do \[ a_{i} \leftarrow \frac{e^{x_{i}-m_{N}}}{d_{N}} \qquad(3)\]

Online softmax

• We want to fuse the operations in a single loop, however (Equation 2) and (Equation 3) cannot be fused as (Equation 3) depends on the value of \(m_N\) which is not known until the end of the loop.

• We can create a surrogate sequence \(d'_{i}\) that is computed in the same way as \(d_i\) but does not depend on \(m_N\): \[ d'_{i} = \sum_{j=1}^i e^{x_j-m_i}. \]

• Furthermore \(d_{N} = d'_{N}\), so we can replace \(d_N\) with \(d'_N\) in (Equation 3).

OnlineSoftmax

• We can find a recurence relation between \(d'_{i}\) and \(d'_{i-1}\): \[ d'_{i} \overset{\text{def}}{=} \sum_{j=1}^i e^{x_j-m_i} \overset{\text{unroll}}{=} \left(\sum_{j=1}^{i-1} e^{x_j-m_i}\right) + e^{x_i-m_i} \\ = \left(\sum_{j=1}^{i-1} e^{x_j-m_i+m_{i-1}-m_{i-1}}\right) + e^{x_i-m_i}\\ = \left(\sum_{j=1}^{i-1} e^{x_j-m_{i-1}}\right)e^{m_{i-1}-{m_{i}}} + e^{x_i-m_i} \\ = d'_{i-1} e^{m_{i-1}-m_i} + e^{x_i-m_i} \]

OnlineSoftmax

2-pass implementation

1. For \(i \leftarrow 1, N\) do \[ m_{i} \leftarrow \max \left(m_{i-1}, x_{i}\right)\\ d'_i \leftarrow d'_{i-1} e^{m_{i-1}-m_i} + e^{x_i-m_i} \qquad(4)\]

2. For \(i \leftarrow 1, N\) do \[ a_{i} \leftarrow \frac{e^{x_{i}-m_{N}}}{d'_{N}} \qquad(5)\]

Multi-pass self-attention

• While we can’t decrease the number of passes for softmax we can decrease the number of passes for the self-attention mechanism, by finding a one-pass recuurrence relation for the O matrix.

• \(Q[k,:]\) is the k-th row vector of the \(Q\) matrix.

• \(K^{T}[:, i]\) : the \(i\)-th column vector of \(K^{T}\) matrix.

• \(O[k,:]\) : the \(k\)-th row of output \(O\) matrix.

• \(V[i,:]\) : the \(i\)-th row of \(V\) matrix.

• \(\left\{\boldsymbol{o}_{i}\right\}: \sum_{j=1}^{i} a_{j} V[j,:]\), a row vector storing partial aggregation result \(A[k,: i] \times V[: i,:]\)

Multi-pass self-attention

1. For \(i \leftarrow 1, N\) do

   \[\begin{aligned} x_{i} & \leftarrow Q[k,:] K^{T}[:, i] \\ m_{i} & \leftarrow \max \left(m_{i-1}, x_{i}\right) \\ d_{i}^{\prime} & \leftarrow d_{i-1}^{\prime} e^{m_{i-1}-m_{i}}+e^{x_{i}-m_{i}} \end{aligned}\]

2. For \(i \leftarrow 1, N\) do \[ a_{i} \leftarrow \frac{e^{x_{i}-m_{N}}}{d_{N}^{\prime}} \qquad(6)\] \[ \boldsymbol{o}_{i} \leftarrow \boldsymbol{o}_{i-1}+a_{i} V[i,:] \qquad(7)\]

3. \[ O[k,:] \leftarrow \boldsymbol{o}_{N} \]

Recurrence relation for \(O\)

• Lets replace the \(a_i\) in (Equation 6) with the recurrence relation from (Equation 7): \[ \boldsymbol{o}_{i}:=\sum_{j=1}^{i}\left(\frac{e^{x_{j}-m_{N}}}{d_{N}^{\prime}} V[j,:]\right) \]

• Depends on \(m_N\) and \(d_N\), so again create a surrogate sequence \(\boldsymbol{o}'_{i}\): \[ \boldsymbol{o}_{i}^{\prime}:=\left(\sum_{j=1}^{i} \frac{e^{x_{j}-m_{i}}}{d_{i}^{\prime}} V[j,:]\right) \]

Recurrence relation for \(O\)

\[\begin{align*} \boldsymbol{o}_{i}^{\prime} & =\sum_{j=1}^{i} \frac{e^{x_{j}-m_{i}}}{d_{i}^{\prime}} V[j,:] \\ & =\left(\sum_{j=1}^{i-1} \frac{e^{x_{j}-m_{i}}}{d_{i}^{\prime}} V[j,:]\right)+\frac{e^{x_{i}-m_{i}}}{d_{i}^{\prime}} V[i,:] \\ & =\left(\sum_{j=1}^{i-1} \frac{e^{x_{j}-m_{i-1}}}{d_{i-1}^{\prime}} \frac{e^{x_{j}-m_{i}}}{e^{x_{j}-m_{i-1}}} \frac{d_{i-1}^{\prime}}{d_{i}^{\prime}} V[j,:]\right)+\frac{e^{x_{i}-m_{i}}}{d_{i}^{\prime}} V[i,:] \\ & =\left(\sum_{j=1}^{i-1} \frac{e^{x_{j}-m_{i-1}}}{d_{i-1}^{\prime}} V[j,:]\right) \frac{d_{i-1}^{\prime}}{d_{i}^{\prime}} e^{m_{i-1}-m_{i}}+\frac{e^{x_{i}-m_{i}}}{d_{i}^{\prime}} V[i,:] \\ \end{align*}\]

Recurrence relation for \(O\)

\[\begin{align*} \boldsymbol{o}_{i}^{\prime} & =\sum_{j=1}^{i} \frac{e^{x_{j}-m_{i}}}{d_{i}^{\prime}} V[j,:] \\ & =\boldsymbol{o}_{i-1}^{\prime} \frac{d_{i-1}^{\prime} e^{m_{i-1}-m_{i}}}{d_{i}^{\prime}}+\frac{e^{x_{i}-m_{i}}}{d_{i}^{\prime}} V[i,:] \end{align*}\]

Flash attention

Single pass self-attention

1. For \(i \leftarrow 1, N\) do \[\begin{aligned} x_i & \leftarrow Q[k,:] K^T[:, i] \\ m_i & \leftarrow \max \left(m_{i-1}, x_i\right) \\ d_i^{\prime} & \leftarrow d_{i-1}^{\prime} e^{m_{i-1}-m_i}+e^{x_i-m_i} \\ \boldsymbol{o}_i^{\prime} & \leftarrow \boldsymbol{o}_{i-1}^{\prime} \frac{d_{i-1}^{\prime} e^{m_{i-1}-m_i}}{d_i^{\prime}}+\frac{e^{x_i-m_i}}{d_i^{\prime}} V[i,:] \end{aligned}\]
2. \[ O[k,:] \leftarrow \boldsymbol{o}_{N}^{\prime}\]

Flash attention 2

• It turns out we don’t need to divide by by \(d_N'\) every time, we can just divide by it once at the end.
• If we look at (Equation 6) we can see we are always dividing \(a_i\) by \(d'_N\), so we can just divide the final result by \(d'_N\).
• We can create a surrogate sequence \(\boldsymbol{s}_{i}\) that is computed in the same way as \(o'_i\) but does not depend on \(d'_N\): \[ \boldsymbol{s}_{i} :=\left(\sum_{j=1}^{i} e^{x_{j}-m_{i}} V[j,:]\right) \]

Flash attention 2

1. For \(i \leftarrow 1, N\) do \[\begin{aligned} x_i & \leftarrow Q[k,:] K^T[:, i] \\ m_i & \leftarrow \max \left(m_{i-1}, x_i\right) \\ d_i^{\prime} & \leftarrow d_{i-1}^{\prime} e^{m_{i-1}-m_i}+e^{x_i-m_i} \\ \boldsymbol{s}_i^{\prime} & \leftarrow \boldsymbol{s}_{i-1}^{\prime} e^{m_{i-1}-m_i}+{e^{x_i-m_i}} V[i,:] \end{aligned}\]
2. \[ O[k,:] \leftarrow \frac{\boldsymbol{s}_{N}^{\prime}}{d_N^{\prime}}\]

GPU programming

CPU vs GPU

• CPU is optimized for low latency and high single-threaded performance, while GPU is optimized for high throughput and parallelism.

GPU programming

Running code on CPU

Running code on CPU (Taken from CS149 Stanford)

GPU programming

Running code on GPU

Running code on GPU(Taken from CS149 Stanford)

CUDA programming

Thread hierarchy

CUDA thread hierarchy (Taken from CS149 Stanford)

CUDA programming

Memory model

CUDA memory model (Taken from CS149 Stanford)

CUDA programming

Memory model

CUDA types of memory (Taken from CS149 Stanford)

CUDA programming

Memory types

CUDA memory types (Taken from yt)

CUDA programming

Memory types

CUDA memory types (Taken from yt)

CUDA programming

Memory types

CUDA memory types [4]

CUDA programming

Synchronization

CUDA synchronization (Taken from CS149 Stanford)

Tiling

Basic matrix multiplication

Tiling

Issues with basic matrix multiplication

Thread memory access [4]

Tiling

Tiled matrix multiplication

Breaking down the matrix into tiles [4]

Tiling

Tiled matmul execution order

Tiled execution order [4]

Tiling

Tiled matrix multiplication kernel

Tiled matmul kernel [4]

Flash Attention

Another graphical representation of Flash Attention [2]

Flash Attention

Algorithm forward pass

Flash attention algorithm [1]

Flash Attention 2

• Flash Attention 2 is very similar, however two tricks are applied:
  • Rescaling is done only once at the end, instead of at every step.
  • No need to keep both maximum \(m^j\) and sum of exponentials \(l^{j}\) instead store \(L^{j}:= m^j + log(l^j)\)

Flash Attention 2

Algorithm forward pass

Flash attention 2 algorithm [2]

Flash Attention 2

• There are some additional improvements added in Flash Attention 2:
  • Worker parralelization is done not only over batch size but also over sequence length.
  • Work paritiionining between warps is changed to reduce memory read/writes.
  • More details in paper [2].

Conclusion

• Flash attention is a highly efficient implementation of the attention mechanism that reduces memory bandwidth usage and computation time.
• It achieves this by reformulating the softmax and attention computations to allow for tiling and single-pass execution.
• Flash attention 2 further improves the efficiency by rescaling only once at the end and optimizing memory access patterns.